import taichi as ti
import math

G = 1  # 重力加速度
PI = 3.1415926  # 圆周率
all_objects = []


@ti.data_oriented
class CelestialObject:

    # 对象构造函数
    # n: 星体个数
    # mass: 星体的质量
    def __init__(self, n, mass):
        self.n = n  # 星体个数
        self.mass = mass  # 星体的质量
        self.pos = ti.Vector.field(2, ti.f32, shape=self.n)  # 2个元素的向量场，表示星体位置
        self.vel = ti.Vector.field(2, ti.f32, shape=self.n)  # 表示星体速度
        self.force = ti.Vector.field(2, ti.f32, shape=self.n)  # 记录星体受力
        all_objects.append(self)

    # 清除星体受力情况
    @ti.kernel
    def clear_force(self):
        for i in self.force:
            # self.force[i] = ti.Vector([0.0, 0.0])
            self.force[i][0] = 0.0
            self.force[i][1] = 0.0

    # 初始化星体的位置，均匀分布到一个圆环上
    # center_x,center_x: 圆环中心点
    # size:
    # init_speed:
    @ti.kernel
    def initialize(self, center_x: ti.f32, center_y: ti.f32, size: ti.f32, init_speed: ti.f32):
        center = ti.Vector([center_x, center_y])
        for i in range(self.n):
            if self.n == 1:
                # 只有一个点的时候放在中心点
                self.pos[i] = ti.Vector([center_x, center_y])
                self.vel[i] = ti.Vector([0.0, 0.0])
            else:
                # 由子类提供generateThetaAndR方法用于计算位置极坐标
                theta, r = self.generateThetaAndR(i, self.n)
                offset_dir = ti.Vector([ti.cos(theta), ti.sin(theta)])
                self.pos[i] = center + r * offset_dir * size
                self.vel[i] = ti.Vector([-offset_dir[1], offset_dir[0]]) * init_speed

    # 计算受力情况
    @ti.kernel
    def compute_force_other(self, other: ti.template()):
        for i in range(self.n):
            p = self.pos[i]
            for j in range(other.n):
                # 计算其他球体对i的万有引力和
                diff = other.pos[j] - p
                r = diff.norm(1e-2)
                self.force[i] += G * (self.mass * other.mass) * diff / r ** 3

    @ti.kernel
    def compute_force_self(self):
        for i in range(self.n):
            p = self.pos[i]
            for j in range(self.n):
                if i == j:
                    continue
                # 计算其他球体对i的万有引力和
                diff = self.pos[j] - p
                r = diff.norm(1e-2)
                f = G * (self.mass * self.mass) * diff / r ** 3
                self.force[i] += f

        # for i in range(self.n):
        #     p = self.pos[i]
        #     for j in range(self.n):
        #         if i < j:
        #             continue
        #         # 计算其他球体对i的万有引力和
        #         diff = self.pos[j] - p
        #         r = diff.norm(1e-2)
        #         f = G * (self.mass * self.mass) * diff / r ** 3
        #         self.force[i] += f
        #         self.force[j] -= f

    # 更新
    # h: 时间系数
    @ti.kernel
    def update_pos(self, h: ti.f32):
        for i in self.vel:
            # 更新速度
            self.vel[i] += h * self.force[i] / self.mass
            # 更新位置
            self.pos[i] += h * self.vel[i]

    def Pos(self):
        return self.pos

    def Mass(self):
        return self.m

    def Number(self):
        return self.n

    def display(self, gui, radius=2, color=0xffffff):
        gui.circles(self.pos.to_numpy(), radius=radius, color=color)


@ti.data_oriented
class Star(CelestialObject):
    def __init__(self, N, mass) -> None:
        super().__init__(N, mass)
        pass

    @staticmethod
    @ti.func
    def generateThetaAndR(i, n):
        theta = 2 * PI * i / ti.cast(n, ti.f32)
        r = 1
        return theta, r


@ti.data_oriented
class Planet(CelestialObject):
    def __init__(self, N, mass) -> None:
        super().__init__(N, mass)
        pass

    @staticmethod
    @ti.func
    def generateThetaAndR(i, n):
        theta = 2 * PI * ti.random()  # theta \in (0, 2PI)
        r = (ti.sqrt(ti.random()) * 0.4 + 0.6)  # r \in (0.6,1)
        return theta, r
